import getopt
import os
import sys
import subprocess


def main(argv):
    list_flag = False
    set_name = ''

    # Read command line args
    try:
        opts, args = getopt.getopt(argv, "ls:", ["set_name="])
    except getopt.GetoptError:
        print("test.py -i <set_name>")
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-l':
            list_flag = True
        if opt in ("-s", "--set_name"):
            set_name = arg

    # BERT args
    input_dir = "swda_data/"
    output_dir = "swda_output/BERT_Features/"

    vocab_file = "BERT_Base/vocab.txt"
    bert_config_file = "BERT_Base/bert_config.json"
    init_checkpoint = "swda_output/BERT_Base/train/model.ckpt-18038"
    layers = -1
    max_seq_length = 128
    batch_size = 8

    # If running on a list of files
    if list_flag:

        # Open file list
        with open(input_dir + set_name + "_split.txt") as file_list:

            # Set input and create output directories
            input_dir += set_name + "_utt/"

            output_dir += set_name + "/"
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, mode=0o777)

            # Run BERT for each file
            for line in file_list:

                input_file = input_dir + line.rstrip() + "_utt.txt"
                output_file = output_dir + line.rstrip() + "_encodings.jsonl"

                command_line_args = "--input_file=" + input_file + " " + \
                                    "--output_file=" + output_file + " " + \
                                    "--vocab_file=" + vocab_file + " " \
                                    "--bert_config_file=" + bert_config_file + " " \
                                    "--init_checkpoint=" + init_checkpoint + " " \
                                    "--layers=" + str(layers) + " " \
                                    "--max_seq_length=" + str(max_seq_length) + " " \
                                    "--batch_size=" + str(batch_size)

                process = subprocess.check_call("nice python extract_features.py " + command_line_args, shell=True)
                
                del process

    # Else just run on the full set
    else:
        input_file = input_dir + set_name + "_set_utt.txt"
        output_file = output_dir + set_name + "_set_encodings.jsonl"

        command_line_args = "--input_file=" + input_file + " " + \
            "--output_file=" + output_file + " " + \
            "--vocab_file=" + vocab_file + " " \
            "--bert_config_file=" + bert_config_file + " " \
            "--init_checkpoint=" + init_checkpoint + " " \
            "--layers=" + str(layers) + " " \
            "--max_seq_length=" + str(max_seq_length) + " " \
            "--batch_size=" + str(batch_size)

        os.system("python extract_features.py " + command_line_args)


if __name__ == "__main__":
    main(sys.argv[1:])
